package com.rtf.framework.common.distributed.ratelimit;

import com.rtf.framework.common.util.UtilString;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.util.Assert;

import java.util.Collections;

/**
 * 分布式限流服务
 * @Author : liupeng
 * @Date : 2020-02-22
 * @Modified By
 */
@Slf4j
public class DistributedRateLimiter {

    private static final String SCRIPT = "local ratelimit_info = redis.pcall('HMGET',KEYS[1],'last_time','current_token') " +
            "local last_time = ratelimit_info[1] " +
            "local current_token = tonumber(ratelimit_info[2]) " +
            "local max_token_seconds = tonumber(ARGV[1]) " +
            "local max_token = tonumber(ARGV[2]) " +
            "local token_rate = tonumber(max_token/max_token_seconds) " +
            "local current_time = tonumber(ARGV[3]) " +
            "local reverse_time = max_token_seconds*1000/max_token " +
            "if current_token == nil then " +
            "  current_token = max_token " +
            "  last_time = current_time " +
            "else " +
            "  local past_time = current_time-last_time " +
            "  local reverse_token = math.floor(past_time/reverse_time) " +
            "  current_token = current_token+reverse_token " +
            "  last_time = reverse_time*reverse_token+last_time " +
            "  if current_token>max_token then " +
            "    current_token = max_token " +
            "  end " +
            "end " +
            "local result = '0' " +
            "if(current_token>0) then " +
            "  result = '1' " +
            "  current_token = current_token-1 " +
            "end  " +
            "redis.call('HMSET',KEYS[1],'last_time',last_time,'current_token',current_token) " +
            "redis.call('pexpire',KEYS[1],math.ceil(reverse_time*(max_token-current_token)+(current_time-last_time))) " +
            "return result";

    private static RedisScript<String> scriptRateLimiter = new DefaultRedisScript<String>(SCRIPT, String.class) ;

    @Autowired
    @Setter
    private StringRedisTemplate stringRedisTemplate ;

    /**
     * 获取限流的key
     * @param code
     * @return
     */
    private String getRateLimiterKey(String code){
        return "ratelimiter:" + code.toLowerCase() ;
    }

    /**
     * 是否能够获取限流锁
     * @param key 锁的key
     * @param maxNum 最大请求量
     * @param timeRange 最大请求量的时间长度，单位秒
     * @return
     */
    public boolean tryAcquire(String key , int timeRange , int maxNum) {
        Assert.isTrue( maxNum>0 , "最大请求量不能为空" );
        Assert.isTrue( timeRange>=1 , "最大请求量的时间长度不能小于1秒" );

        String flagKey = getRateLimiterKey( key ) ;
        // 执行lua脚本，第一个参数：最大并发量，第二个参数：每秒并发量 第三个参数：当前时间
        String result = stringRedisTemplate.execute( scriptRateLimiter ,
                stringRedisTemplate.getStringSerializer() , stringRedisTemplate.getStringSerializer() ,
                Collections.singletonList(flagKey) , timeRange+"" , maxNum+"" , System.currentTimeMillis()+"" ) ;

        return UtilString.equalsIgnoreCase( result , "1" );
    }

    /**
     * 释放redis锁
     * @param key
     */
    public void release(String key) {
        String flagKey = getRateLimiterKey( key ) ;

        stringRedisTemplate.delete( flagKey ) ;
    }

//    protected static final RedisStandaloneConfiguration getStandaloneConfig(RedisProperties redisProperties) {
//        RedisStandaloneConfiguration config = new RedisStandaloneConfiguration();
//        config.setHostName(redisProperties.getHost());
//        config.setPort(redisProperties.getPort());
//        config.setPassword(RedisPassword.of(redisProperties.getPassword()));
//        config.setDatabase(redisProperties.getDatabase());
//        return config;
//    }
//
//    public static void main(String[] args) throws Exception {
//        RedisProperties redisProperties = new RedisProperties() ;
//        redisProperties.setHost("10.254.23.150");
//        redisProperties.setPort(18000);
//
//
//        LettuceClientConfiguration lettuceClientConfiguration = LettuceClientConfiguration.builder().build() ;
//
//        AppRedisHostListDefault appRedisHostList = new AppRedisHostListDefault() ;
//        appRedisHostList.setProperties( redisProperties ) ;
//
//        AppCodisConnectionFactory redisConnectionFactory = new AppCodisConnectionFactory( getStandaloneConfig(redisProperties) ,
//                lettuceClientConfiguration ) ;
//        redisConnectionFactory.afterPropertiesSet() ;
//
//        StringRedisTemplate stringRedisTemplate = new StringRedisTemplate( redisConnectionFactory );
//        stringRedisTemplate.afterPropertiesSet();
//
//
//        DistributedRateLimiter distributedRateLimiter = new DistributedRateLimiter() ;
//        distributedRateLimiter.setStringRedisTemplate( stringRedisTemplate ) ;
//
//        while (true){
//            boolean got = distributedRateLimiter.tryAcquire("test" , 10 , 1) ;
//
//            log.info("{} {}" , new DateTime().toString("HH:mm:ss.SSS") , got ) ;
//
//            Thread.currentThread().sleep(1000);
//        }
//
//    }

}
